import numpy as np
from torch.utils.data import Dataset
from torch.nn import Module, Linear, LSTM, Dropout, ReLU, init

class TimeSeriesDataset(Dataset):
	def __init__(self, x, y):
		self.x = x.astype(np.float32)
		self.y = y.astype(np.float32)
		
	def __len__(self):
		return len(self.x)

	def __getitem__(self, idx):
		return (self.x[idx], self.y[idx])

class LSTMModel(Module):
	def __init__(self, input_size=1, hidden_layer_size=32, num_layers=2, output_size=1, dropout=0.2):
		super().__init__()
		self.hidden_layer_size = hidden_layer_size

		self.linear_1 = Linear(input_size, hidden_layer_size)
		self.relu = ReLU()
		self.lstm = LSTM(hidden_layer_size, hidden_size=self.hidden_layer_size, num_layers=num_layers, batch_first=True)
		self.dropout = Dropout(dropout)
		self.linear_2 = Linear(num_layers*hidden_layer_size, output_size)
		
		self.init_weights()

	def init_weights(self):
		for name, param in self.lstm.named_parameters():
			if 'bias' in name:
				init.constant_(param, 0.0)
			elif 'weight_ih' in name:
				init.kaiming_normal_(param)
			elif 'weight_hh' in name:
				init.orthogonal_(param)

	def forward(self, x):
		batchsize = x.shape[0]

		# layer 1
		x = self.linear_1(x)
		x = self.relu(x)
		
		# LSTM layer
		lstm_out, (h_n, c_n) = self.lstm(x)

		# reshape output from hidden cell into [batch, features] for `linear_2`
		x = h_n.permute(1, 0, 2).reshape(batchsize, -1) 
		
		# layer 2
		x = self.dropout(x)
		predictions = self.linear_2(x)
		return predictions
